from langchain_core.prompts import ChatPromptTemplate
from langchain_core.tools import Tool

from common.llm_builder import llm_messages, llm


# 定义工具方法
def add(a: int, b: int) -> int:
    return a + b


# 1-创建一个ChatPromptTemplate
chat_prompt_template = ChatPromptTemplate.from_messages(llm_messages)
print("-----chat_prompt_template_type:", type(chat_prompt_template))
print("-----llm_type:", type(llm))

# 2-构建工具函数
add_tool = Tool.from_function(
    func=add,
    name="add",
    description="计算两个数字的和"
)
llm_bind_tool = llm.bind_tools(tools=[add_tool])
print("-----llm_bind_tool:", llm_bind_tool)
print("-----llm_bind_tool_type:", type(llm_bind_tool))

# 3-构建链式对象
llm_chain = chat_prompt_template | llm_bind_tool

# 4-打印响应结果
chat_prompt_response = llm_chain.invoke({
    "role": "数学",
    "domain": "数字计算",
    "question": "1+150？"
})

# 4-打印响应结果
print("-----Hold on, LLM 正在回答！-----")
print(chat_prompt_response)
